import math
import map_decrete as md
# 2024年04月18日23:02:06 实现最基础的A星功能
# 输入输出
"""
A star planning algorithm
input: map grid with 0 and 1, represents barrier and passable brick
        start point and end point
output: path from start to end
        the number of visited times of each node

"""


# 伪代码
"""
fake code:
function AStar(start, goal)
    openList := {start}
    closedList := {}

    while openList is not empty
        current := node in openList with the lowest f_cost

        if current is goal
            return path

        move current from openList to closedList

        for each neighbor of current
            if neighbor is not traversable or neighbor is in closedList
                skip to the next neighbor

            if neighbor is not in openList
                add neighbor to openList
                set neighbor's parent to current
                calculate neighbor's g_cost, h_cost, and f_cost

            else if neighbor is in openList and current's g_cost is lower than neighbor's g_cost
                update neighbor's parent and g_cost

    return no path found

"""


class Node:
    def __init__(self, x, y, traversable=True):
        self.x = x
        self.y = y
        self.traversable = traversable
        self.g_cost = 0
        self.h_cost = 0
        self.f_cost = 0
        self.parent = [0, 0]    # 使用列表存储父节点坐标，可以修改
        self.visited_times = 0


def create_node_matrix(map_grid):
    height, width = map_grid.shape
    nodes = [[Node(x, y, map_grid[x][y]) for y in range(width)] for x in range(height)]
    return nodes


def calc_f_cost(node):
    node.f_cost = node.g_cost + node.h_cost
    return node.f_cost


class AstarPlanner:
    def __init__(self, nodes, start, end):      # start 用元组传递
        self.node_matrix = nodes
        self.xmax = len(nodes)
        self.ymax = len(nodes[0])
        self.start = start
        self.end = end
        self.height = len(nodes)
        self.width = len(nodes[0])
        self.open_set = dict()  # 元素是键值对 (x, y) : f_cost   list is unhashable
        self.closed_set = dict()  # 元素是键值对 (x, y) : f_cost
        self.path = []

    # 使用欧式距离计算h_cost
    def calc_h_cost(self, node):
        if not node.h_cost:
            node.h_cost = math.sqrt((node.x - self.end[0]) ** 2 + (node.y - self.end[1]) ** 2)
        return node.h_cost

    def add_to_open_set(self, node, parent_node):
        x, y = node.x, node.y
        this_node = self.node_matrix[x][y]
        this_node.h_cost = self.calc_h_cost(this_node)
        this_node.g_cost = parent_node.g_cost + 1
        this_node.f_cost = calc_f_cost(this_node)
        this_node.parent = [parent_node.x, parent_node.y]
        self.open_set.update({(x, y): this_node.f_cost})

    def get_path(self):
        path = [self.end]
        while path[-1] != self.start:
            parent = self.node_matrix[path[-1][0]][path[-1][1]].parent
            path.append(tuple(parent))
        return path
        # TODO 20240320 to finish this method

    def get_neighbours(self, x, y):
        if 0 < x < self.xmax-1:
            x_range = [x-1, x, x+1]
        elif x == 0:
            x_range = [0, 1]
        else:
            x_range = [self.xmax-2, self.xmax-1]

        if 0 < y < self.ymax-1:
            y_range = [y-1, y, y+1]
        elif y == 0:
            y_range = [0, 1]
        else:
            y_range = [self.ymax-2, self.ymax-1]

        neighbours = [(i, j) for i in x_range for j in y_range if (i, j) != (x, y)]
        return neighbours

    def planning(self):
        start_node = self.node_matrix[self.start[0]][self.start[1]]
        end_node = self.node_matrix[self.end[0]][self.end[1]]
        start_node.h_cost = self.calc_h_cost(start_node)
        start_node.g_cost = 0
        start_node.f_cost = start_node.h_cost + start_node.g_cost
        self.open_set.update({(self.start[0], self.start[1]): start_node.f_cost})

        while self.open_set:
            [x, y] = min(self.open_set, key=self.open_set.get)
            current_node = self.node_matrix[x][y]
            current_node.visited_times += 1

            if (x, y) == self.end:
                path = self.get_path()
                return path

            else:
                self.closed_set.update({(x, y): current_node.f_cost})
                self.open_set.pop((x, y))

                neighbours = self.get_neighbours(x, y)
                for neighbour in neighbours:
                    neighbour_node = self.node_matrix[neighbour[0]][neighbour[1]]
                    if not neighbour_node.traversable or neighbour in self.closed_set:
                        continue

                    if neighbour not in self.open_set:
                        self.add_to_open_set(neighbour_node, current_node)

                    elif neighbour in self.open_set and current_node.g_cost < neighbour_node.g_cost:
                        neighbour_node.parent = [current_node.x, current_node.y]
                        neighbour_node.g_cost = current_node.g_cost + 1
                        neighbour_node.f_cost = calc_f_cost(neighbour_node)
                        self.open_set.update({neighbour: neighbour_node.f_cost})

        return []   # no path found


def main():
    new_map = md.map_resize(md.map_grid, 50, 50)
    create_node_matrix(new_map)
    start, end = (0, 0), (49, 49)
    astar = AstarPlanner(create_node_matrix(new_map), start, end)
    path = astar.planning()


if __name__ == '__main__':
    main()
